from sklearn.mixture import GaussianMixture
from sklearn.model_selection import train_test_split
import numpy as np
import joblib
from .utils.logger import setup_logger

logger = setup_logger("GMMTrainer")

class GMMTrainer:
    def __init__(self, config: dict):
        self.config = config
        self.model = None

    def _select_optimal_k(self, X: np.ndarray) -> int:
        """通过BIC选择最佳聚类数"""
        bic_scores = []
        k_range = range(
            self.config['bic_range'][0],
            self.config['bic_range'][1]+1
        )

        logger.info(f"Evaluating BIC for k in {k_range}")
        for k in k_range:
            gmm = GaussianMixture(
                n_components=k,
                covariance_type=self.config['covariance_type'],
                max_iter=self.config['max_iter'],
                tol=self.config['tolerance'],
                init_params='kmeans'
            )
            gmm.fit(X)
            bic_scores.append(gmm.bic(X))

        return k_range[np.argmin(bic_scores)]

    def train(self, X: np.ndarray) -> GaussianMixture:
        """完整训练流程"""
        logger.info("Starting GMM training pipeline")

        # 数据集划分（用于后续验证）
        X_train, X_val = train_test_split(X, test_size=0.2, random_state=42)

        # 自动选择K值
        optimal_k = self._select_optimal_k(X_train)
        logger.info(f"Selected optimal k: {optimal_k}")

        # 训练最终模型
        self.model = GaussianMixture(
            n_components=optimal_k,
            covariance_type=self.config['covariance_type'],
            max_iter=self.config['max_iter'],
            tol=self.config['tolerance'],
            init_params='kmeans'
        )
        self.model.fit(X_train)

        # 模型验证
        train_score = self.model.score(X_train)
        val_score = self.model.score(X_val)
        logger.info(f"Model trained | Train LL: {train_score:.2f} | Val LL: {val_score:.2f}")

        return self.model